#include <stdio.h>
#include <string.h>
#include <stdint.h>

// Enable ECB, CTR and CBC mode. Note this can be done before including aes.h or at compile-time.
// E.g. with GCC by using the -D flag: gcc -c aes.c -DCBC=0 -DCTR=1 -DECB=1
#define CBC 1
#define CTR 0
#define ECB 0

#include "aes.h"

static void phex(uint8_t* str);
static int test_encrypt_cbc(void);
static int test_decrypt_cbc(void);

int main(void)
{
    int exit;
    printf("\nTesting AES128\n\n");

    exit = test_encrypt_cbc() + test_decrypt_cbc();
    return exit;
}

// prints string as hex
static void phex(uint8_t* str)
{

#if defined(AES256)
    uint8_t len = 32;
#elif defined(AES192)
    uint8_t len = 24;
#elif defined(AES128)
    uint8_t len = 16;
#endif

    unsigned char i;
    for (i = 0; i < len; ++i)
        printf("%.2x", str[i]);
    printf("\n");
}

static int test_encrypt_cbc(void)
{
    uint8_t key[] = { 0x62, 0x69, 0x61, 0x6E, 0x78, 0x69, 0x65, 0x32, 0x32, 0x30, 0x32, 0x5F, 0x32, 0x32, 0x30, 0x33 };
    uint8_t in[]  = { 0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b };

    uint8_t out[] = { 0x18, 0xca, 0x09, 0x0c, 0x19, 0xe5, 0xf3, 0xb4, 0x1d, 0x19, 0xfa, 0xf5, 0x31, 0x41, 0xb5, 0x0a };
    uint8_t iv[]  = { 0x66, 0x72, 0x61, 0x6D, 0x65, 0x64, 0x61, 0x74, 0x61, 0x70, 0x61, 0x72, 0x74, 0x61, 0x65, 0x73 };


    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CBC_encrypt_buffer(&ctx, in, 16);

    printf("CBC encrypt: ");

    if (0 == memcmp((char*) out, (char*) in, 16)) {
        printf("SUCCESS!\n");
	return(0);
    } else {
        printf("FAILURE!\n");
	return(1);
    }
}

static int test_decrypt_cbc(void)
{
    uint8_t key[] = { 0x62, 0x69, 0x61, 0x6E, 0x78, 0x69, 0x65, 0x32, 0x32, 0x30, 0x32, 0x5F, 0x32, 0x32, 0x30, 0x33 };
    uint8_t iv[]  = { 0x66, 0x72, 0x61, 0x6D, 0x65, 0x64, 0x61, 0x74, 0x61, 0x70, 0x61, 0x72, 0x74, 0x61, 0x65, 0x73 };

    uint8_t in[16] = { 0x18, 0xca, 0x09, 0x0c, 0x19, 0xe5, 0xf3, 0xb4, 0x1d, 0x19, 0xfa, 0xf5, 0x31, 0x41, 0xb5, 0x0a };
    uint8_t out[16]  = { 0x68, 0x65, 0x6c, 0x6c, 0x6f };

    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CBC_decrypt_buffer(&ctx, in, 16);

    printf("CBC decrypt: ");

    if (0 == memcmp((char*) out, (char*) in, 5)) {
        printf("SUCCESS!\n");
	return(0);
    } else {
        printf("FAILURE!\n");
	return(1);
    }
}
